import dgl
import torch
import numpy as np
import json
import torch.nn as nn
import torch.nn.functional as F
import dgl.data
import networkx as nx
from collections import Counter


def build_graph_func(data_file_path,family_file_path,unknown_family_str):
    # Load the JSON file to examine its structure

    with open(data_file_path, 'r') as file:
        data = json.load(file)
    with open(family_file_path, 'r') as file:
        family_file = json.load(file)

    # Step 0:
    family_count = Counter(family_file.values())
    sorted_family_count = sorted(family_count.items(), key=lambda x: x[1], reverse=True)
    sorted_family_keys = [item[0] for item in sorted_family_count] # 0:136 共137个家族


    # Step 1: Create a list of unique samples and system calls
    sample_names = list(data.keys()) #样本名称列表

    # 创建一个空集合来存储系统调用名称和系统调用号
    system_calls = set() 
    # 遍历 JSON 数据的第一级键（样本名称）
    for sample in data.values():
        # 然后遍历每个样本的第二级键（系统调用名称和号码）
        for sys_call_name, sys_call_number in sample.items():
            # 将系统调用名称和号码作为元组添加到集合中
            system_calls.add((sys_call_name, sys_call_number))
    system_calls = list(system_calls) #系统调用名称列表

    # Step 2: Create mappings for samples and system calls to assign unique ids
    # 构建样本节点的map
    sample_idx_map = {}
    unique_id = 0  # Starting ID

    for sample_name in sample_names:
        sample_idx_map[sample_name] = unique_id
        unique_id += 1

    # 构建系统调用节点名称的map
    # 格式为{'recvmmsg': 781, 'prlimit64': 782}
    sys_call_name_idx_map = {}
    unique_id = len(sample_names)  # Resetting ID for system calls

    for sys_call in system_calls:
        sys_call_name_idx_map[sys_call[0]] = unique_id
        unique_id += 1

    # 构建系统调用id号的map
    # 格式为{299: 781, 302: 782}
    sys_call_id_idx_map = {}
    unique_id = len(sample_names)  # Resetting ID for system calls

    for sys_call in system_calls:
        sys_call_id_idx_map[sys_call[1]] = unique_id
        unique_id += 1

    # Step 3: Prepare edge data
    src_edges = []  # Source nodes (samples)
    dst_edges = []  # Destination nodes (system calls)

    for sample, syscalls in data.items():
        sample_node_id = sample_idx_map[sample]
        for syscall in syscalls.keys():
            syscall_node_id = sys_call_name_idx_map[syscall]
            src_edges.append(sample_node_id)
            dst_edges.append(syscall_node_id)

    # Convert edge lists to tensors
    src_edges = torch.tensor(src_edges, dtype=torch.int64)
    dst_edges = torch.tensor(dst_edges, dtype=torch.int64)

    # Step 4: Create the DGL graph
    # g = dgl.graph((src_edges, dst_edges))
    g = dgl.graph((dst_edges, src_edges))

    # Step 5: 给定节点类型/标签
    label = torch.zeros(g.number_of_nodes(), dtype=torch.int32)
    g.ndata['label'] = label

    g.ndata['malware_family'] = torch.zeros(g.number_of_nodes(), dtype=torch.int32)

    # malware_families_list = []

    for sample_id in list(sample_idx_map.values()):
        if "VirusShare" in list(sample_idx_map.items())[sample_id][0]:
            # 恶意样本
            g.ndata['malware_family'][sample_id] = sorted_family_keys.index(family_file.get(list(sample_idx_map.items())[sample_id][0][11:-8]))
            g.ndata['label'][sample_id] = 1
        else:
            # 良性样本
            g.ndata['label'][sample_id] = 0
            g.ndata['malware_family'][sample_id] = 14
    

    # exit()
    for syscall_idx in list(sys_call_name_idx_map.values()):
        # 系统调用节点
        g.ndata['label'][syscall_idx] = 2
        g.ndata['malware_family'][syscall_idx] = 15

    # Step 6: 给定节点特征
    # 初始化一个全为零的特征张量，长度等于图中节点的数量
    features = torch.zeros(g.number_of_nodes(), 335, dtype=torch.float32)
    # 遍历每个系统调用，为其分配特征（系统调用号）
    for sys_call_name, sys_call_number in system_calls:
        # 获取系统调用节点的索引
        node_idx = sys_call_name_idx_map[sys_call_name]
        # 将整数转换为 tensor
        sys_call_number_tensor = torch.tensor([sys_call_number], dtype=torch.int64)
        # 使用 one_hot 函数创建 one-hot 编码
        one_hot_vector = F.one_hot(sys_call_number_tensor, num_classes=335)

        # 将系统调用号设置为该节点的特征
        features[node_idx] = one_hot_vector

    # 将特征张量分配给图的数据
    g.ndata['feat'] = features

    # Step 7: 数据集划分

    # 获取所有标签为0和1的节点的索引
    labels = g.ndata['label']
    benign_nodes = torch.where(labels == 0)[0] # 581个
    malicious_nodes = torch.where(labels == 1)[0] # 1000个
    # sys_call_nodes = torch.where(labels == 2)[0] #1581:1827
    # print(g.ndata['feat'][1581:1827].max()) # 最大值334

    # 初始化训练和测试掩码
    train_mask = torch.zeros_like(labels, dtype=torch.bool)
    test_mask = torch.zeros_like(labels, dtype=torch.bool)

    # 定义划分比例
    train_ratio = 0.80  # 80% for training

    # 为良性样本（label 0）划分训练和测试集
    benign_nodes = benign_nodes.numpy()  # Convert to numpy array
    np.random.shuffle(benign_nodes)
    benign_nodes = torch.from_numpy(benign_nodes)  # Convert back to tensor
    num_train_benign = int(len(benign_nodes) * train_ratio)
    train_mask[benign_nodes[:num_train_benign]] = True
    test_mask[benign_nodes[num_train_benign:]] = True

    # 为恶意样本（label 1）划分训练和测试集
    #g.ndata['malware_family'] 是每个节点的恶意软件家族标签
    malware_families = g.ndata['malware_family']
    # print(malware_families.unique())
    # exit()
    # 选择一个家族作为测试集的一部分
    # 例如，选择家族ID为2:136的家族作为测试集
    # test_family_ids = list(range(2, 136))
    test_family_str = unknown_family_str
    test_family_ids = sorted_family_keys.index(test_family_str)
    # print(test_family_ids)
    # exit()
    # test_family_nodes = torch.where(malware_families.isin(test_family_ids))[0]
    # 使用列表推导式和any函数来检查家族ID
    test_family_nodes = torch.tensor([idx for idx, family in enumerate(malware_families) if family.item() == test_family_ids])
    # print(test_family_nodes.numel())
    # exit()
    # 更新恶意节点列表，排除用于测试的家族
    malicious_nodes = torch.tensor([idx for idx, family in enumerate(malware_families) if labels[idx] == 1 and family.item() != test_family_ids])

    # 更新恶意节点列表，排除用于测试的家族
    # malicious_nodes = torch.where((labels == 1) & (~malware_families.isin(test_family_ids)))[0]

    # 划分剩余恶意样本
    malicious_nodes = malicious_nodes.numpy()  # Convert to numpy array
    np.random.shuffle(malicious_nodes)
    malicious_nodes = torch.from_numpy(malicious_nodes)  # Convert back to tensor
    num_train_malicious = int(len(malicious_nodes) * train_ratio)
    train_mask[malicious_nodes[:num_train_malicious]] = True
    test_mask[malicious_nodes[num_train_malicious:]] = True

    # 对于测试家族的所有样本，将它们添加到测试集中
    test_mask[test_family_nodes] = True

    # 将掩码赋值给图的ndata
    g.ndata['train_mask'] = train_mask
    g.ndata['test_mask'] = test_mask

    # nx_graph = dgl.to_networkx(g)

    # nx.write_gexf(nx_graph, "node_class_graph.gexf")
    # exit()
    return g,sorted_family_keys

